import numpy as np


# 加载数据集
def load_data():
    with open('../data/决策树数据.txt') as fr:
        lines = fr.readlines()
    x = np.empty((len(lines), 7), dtype=int)
    for i in range(len(lines)):
        line = lines[i].strip().split(',')
        x[i] = line
    test_x = x[10:]
    x = x[:10]
    return x, test_x


x, test_x = load_data()


# print(x.shape, test_x.shape)

# 计算数据集的熵
def get_entropy(_x):
    entropy = 0
    # 统计y的熵
    y = _x[:, -1]
    # 统计每个结果出现的次数,[5,5],表示0出现了5次，1出现了5次
    bincount = np.bincount(y)
    for count in bincount:
        if count == 0:
            continue

        # 出现次数/总次数=出现概率
        prob = count / len(_x)

        # 熵 = p * log(p) * -1
        entropy -= prob * np.log2(prob)

    return entropy


# print(get_entropy(x))#1.0 数据的混乱程度达到最大

# 计算信息增益
def get_gain(_x, col):
    # 列熵
    col_entropy = 0

    # 根据列的值，把数据分成n份
    for val in set(_x[:, col]):
        x_by_col_and_value = _x[_x[:, col] == val]
        # 这个数据子集的概率
        prob = len(x_by_col_and_value) / len(_x)
        # 求这个数据子集的熵
        entropy = get_entropy(x_by_col_and_value)
        # 这个列的熵，等于这个式子的累积
        col_entropy += prob * entropy
    # 信息增益，就是切分数据后，熵值能下降多少，这个值越大越好
    gain = get_entropy(_x) - col_entropy
    # 用这个就是id3决策树，他倾向于选择可取值多的列
    return gain


print(get_gain(x, 0))
